import torch

# 初始化：输入 x 和参数 A（需要梯度）
x = torch.tensor([1.0, 2.0], requires_grad=False)
A = torch.nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))  # 可学习参数

# 定义优化器
optimizer = torch.optim.SGD([A], lr=0.1)

# 清除旧梯度
print('---------------------')
optimizer.zero_grad()
print("Before backward: A.grad =\n", A.grad)

# 前向传播：构造一个简单的 loss
print('---------------------')
y = A @ x            # 矩阵乘法
loss = y.sum()       # 简单的 loss
print("Loss:", loss.item())

# 反向传播：计算 A 的梯度
print('---------------------')
loss.backward()
print("After backward: A.grad =\n", A.grad)

# 更新参数：使用 A.grad 更新 A.data
print('---------------------')
optimizer.step()
print("After step: A =\n", A.data)


print('---------------------')
print("Before backward: A.grad =\n", A.grad)
optimizer.zero_grad()
print("Before backward: A.grad =\n", A.grad)
